import numpy as np
from sklearn.neighbors import BallTree



# 1. 创建示例数据集（三维空间中的点）
# 生成20个随机三维数据点 (20, 3)
np.random.seed(42)
data = np.random.rand(20, 3)
print("示例数据集 (前5个点):\n", data[:5])

# 2. 构建BallTree索引（使用欧氏距离）
tree = BallTree(data, leaf_size=2, metric='euclidean')

# 3. 创建查询点
query_point = np.array([[0.5, 0.5, 0.5]])  # 三维查询点

# 4. 执行k近邻搜索 (k=3)
k = 3
distances, indices = tree.query(query_point, k=k)
print(f"\n查询点: {query_point[0]}")
print(f"最近的{k}个邻居的索引: {indices[0]}")
print(f"距离: {distances[0]}")

# 5. 执行半径搜索 (半径=0.5)
radius = 0.5
indices_radius = tree.query_radius(query_point, r=radius)[0]
print(f"\n半径{radius}内的邻居索引: {indices_radius}")

# 可选：获取半径内的具体距离
if len(indices_radius) > 0:
    distances_radius = tree.query(query_point, k=len(indices_radius), return_distance=True)[0][0]
    print(f"距离: {distances_radius}")
else:
    print("半径内未找到邻居")

# 6. 可视化结果（仅适用于2D/3D数据）
try:
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

    fig = plt.figure(figsize=(10, 7))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制所有数据点
    ax.scatter(data[:,0], data[:,1], data[:,2], c='blue', s=50, label='数据点')
    
    # 绘制查询点
    ax.scatter(query_point[0,0], query_point[0,1], query_point[0,2], 
               c='red', s=100, marker='*', label='查询点')
    
    # 绘制k近邻点
    neighbor_points = data[indices[0]]
    ax.scatter(neighbor_points[:,0], neighbor_points[:,1], neighbor_points[:,2],
               c='green', s=80, edgecolors='black', label=f'Top-{k}邻居')
    
    # 绘制半径内邻居
    if len(indices_radius) > 0:
        radius_points = data[indices_radius]
        ax.scatter(radius_points[:,0], radius_points[:,1], radius_points[:,2],
                   c='none', s=100, edgecolors='purple', linewidths=2, label='半径内邻居')
    
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title('BallTree 最近邻搜索演示')
    ax.legend()
    plt.show()

except ImportError:
    print("\n未安装matplotlib，跳过可视化部分")